from langchain_community.utilities import SQLDatabase
from langchain_ollama import OllamaLLM # 使用Ollama本地部署的Llama3.2
from langchain.agents import AgentType, initialize_agent
from langchain.tools import Tool
from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate
import sqlite3
"""
解决：
(1) 报错：ValueError: This function requires a bind_tools() method be implemented on the LLM.意思是：因为create_sql_agent需要LLM支持工具调用功能引起的。
（3）解决二：使用简单agent，自定义SQL工具。
"""

# 1. 初始化SQLite数据库
def setup_database():
    """设置示例数据库"""
    conn = sqlite3.connect('example.db')
    cursor = conn.cursor()

    # 创建用户表
    cursor.execute('''
    CREATE TABLE IF NOT EXISTS users (
        id INTEGER PRIMARY KEY,
        name TEXT NOT NULL,
        age INTEGER,
        email TEXT
    )
    ''')

    # 清空现有数据并插入新数据
    cursor.execute("DELETE FROM users")
    cursor.execute("INSERT INTO users (name, age, email) VALUES ('Alice', 30, 'alice@email.com')")
    cursor.execute("INSERT INTO users (name, age, email) VALUES ('Bob', 25, 'bob@email.com')")
    cursor.execute("INSERT INTO users (name, age, email) VALUES ('Charlie', 35, 'charlie@email.com')")

    conn.commit()
    conn.close()


# 初始化数据库
setup_database()

# 2. 连接到数据库
db = SQLDatabase.from_uri("sqlite:///example.db")

# 3. 初始化Llama3.2模型
llm = OllamaLLM(model="llama3.2:3b")


# 4. 创建自定义SQL工具
def run_sql_query(query):
    """执行SQL查询并返回结果"""
    try:
        result = db.run(query)
        return result
    except Exception as e:
        return f"查询错误: {str(e)}"


# 创建工具
sql_tool = Tool(
    name="SQLQuery",
    func=run_sql_query,
    description="用于执行SQL查询的工具。输入应该是有效的SQL查询语句。"
)

# 5. 创建简单的代理
tools = [sql_tool]

# 使用更简单的代理类型
agent = initialize_agent(
    tools=tools,
    llm=llm,
    agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,  # 使用零样本代理
    verbose=True,
    handle_parsing_errors=True  # 处理解析错误
)

# 6. 创建提示模板来帮助模型生成SQL
sql_prompt = PromptTemplate(
    input_variables=["question"],
    template="""
    你是一个SQL专家。请将以下自然语言问题转换为SQL查询语句。
    数据库结构：users表包含id, name, age, email字段。

    问题：{question}

    只返回SQL查询语句，不要有其他解释。
    SQL查询："""
)

sql_chain = LLMChain(llm=llm, prompt=sql_prompt)

# 7. 执行查询的函数
def query_database(question):
    """处理自然语言查询"""
    try:
        # 第一步：将问题转换为SQL
        sql_query = sql_chain.run(question)
        print(f"生成的SQL: {sql_query}")

        # 第二步：执行SQL查询
        result = run_sql_query(sql_query)
        print(f"查询结果: {result}")

        # 第三步：用自然语言解释结果
        explanation_prompt = f"""
        根据以下SQL查询结果，用简单的中文解释：

        问题：{question}
        SQL查询：{sql_query}
        结果：{result}

        请用自然语言解释这个结果：
        """

        explanation = llm.invoke(explanation_prompt)
        return explanation

    except Exception as e:
        return f"处理查询时出错: {str(e)}"


# 8. 测试查询
if __name__ == "__main__":
    # 示例查询
    questions = [
        "查询所有用户的姓名和年龄",
        "年龄大于28岁的用户有哪些？",
        "统计用户的总数"
    ]

    for question in questions:
        print(f"\n{'=' * 50}")
        print(f"问题: {question}")
        print(f"{'=' * 50}")

        result = query_database(question)
        print(f"答案: {result}")
        print(f"{'=' * 50}")